package cn.cdutetc.foxwatch.ws;

import cn.cdutetc.foxwatch.config.GetHttpSessionConfig;
import com.alibaba.dashscope.aigc.multimodalconversation.*;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.exception.UploadFileException;
import com.alibaba.fastjson2.JSON;
import io.reactivex.Flowable;
import jakarta.servlet.http.HttpSession;
import jakarta.websocket.*;
import jakarta.websocket.server.ServerEndpoint;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author 王勇
 * @ClassName： test
 * @ClassPath： cn.cdutetc.foxwatch.ws.test
 * @Description：ai聊天接口 ws://localhost:8089/chat/ai
 * @date 2024/10/3 00:20
 */

@Component
@ServerEndpoint(value = "/chat/ai", configurator = GetHttpSessionConfig.class)
public class AiEndpoint {

    private HttpSession httpSession;

    private static String getSessionUserId(Session session) {
        List<String> useList = session.getRequestParameterMap().get("userId");
        if (useList != null) {
            return useList.get(0);
        } else {
            return null;
        }
    }

    @OnOpen
    public void onOpen(Session session, EndpointConfig config) {
        System.out.println("WebSocket opened");
    }

    @OnMessage
    public void onMessage(String message, Session session) throws IOException, NoApiKeyException, UploadFileException {
        System.out.println(message);
        Message msg = JSON.parseObject(message, Message.class);
        // 初始化对话
        MultiModalConversation conv = new MultiModalConversation();
        MultiModalMessageItemText systemText = new MultiModalMessageItemText("You are an AI Psychological counselor, and your name is 心理咨询师. Remember your identity, you can only answer questions related to psychology.");
        MultiModalConversationMessage systemMessage = MultiModalConversationMessage.builder()
                .role(Role.SYSTEM.getValue()).content(Arrays.asList(systemText)).build();
        MultiModalMessageItemText userText = new MultiModalMessageItemText(msg.getMessage());
        MultiModalConversationMessage userMessage = MultiModalConversationMessage.builder()
                .role(Role.USER.getValue()).content(Arrays.asList(userText)).build();
        List<MultiModalConversationMessage> messages = new ArrayList<>();
        messages.add(systemMessage);
        messages.add(userMessage);
        MultiModalConversationParam param = MultiModalConversationParam.builder()
                .model(MultiModalConversation.Models.QWEN_VL_CHAT_V1)
                .messages(messages)
                .build();
        // 获取数据流
        Flowable<MultiModalConversationResult> result = conv.streamCall(param);
        result.blockingForEach(item -> {
            System.out.println(item);
            try {
                session.getBasicRemote().sendText(JSON.toJSONString(item.getOutput().getChoices().get(0).getMessage().getContent().get(0).get("text").toString()));
            } catch (IOException e) {
                // 处理异常情况
                e.printStackTrace();
            }
        });
        session.getBasicRemote().sendText("end");
    }

    @OnClose
    public void onClose(Session session) {

    }

    @OnError
    public void onError(Session session, Throwable throwable) {
        System.out.println("WebSocket error: " + session.getId());
        throwable.printStackTrace();
    }
}